from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

if __name__ == "__main__":

    corrs = []
    projection_dim = 1
    sample_dim = 30
    n_samples = 2000

    projection_dim_range = sample_dim*10

    for pd in range(1,int(projection_dim_range)):

        progress_bar(pd, projection_dim_range)

        # 生成随机矩阵，元素值在 [-1,1] 之间
        random_matrix = np.random.rand(sample_dim, pd) * 2 - 1

        # 生成随机向量，元素值在 [-1,1] 之间
        random_vectors = np.random.rand(n_samples, sample_dim) * 2 - 1

        # 使用 random_matrix 将 random_vectors 投影
        random_vectors_prj = np.dot(random_vectors, random_matrix)

        print("shape of random_vectors_prj: ", random_vectors_prj.shape)

        # 计算 random_vectors 的距离矩阵
        random_vectors_dist_mat = squareform(pdist(random_vectors, 'euclidean'))

        # 计算 random_vectors_prj 的距离矩阵
        random_vectors_prj_dist_mat = squareform(pdist(random_vectors_prj, 'euclidean'))

        # 将 random_vectors_dist_mat 和 random_vectors_prj_dist_mat 展平为向量
        random_vectors_dist_mat_vector = random_vectors_dist_mat.reshape(-1)
        random_vectors_prj_dist_mat_vector = random_vectors_prj_dist_mat.reshape(-1)

        # 计算 random_vectors_dist_mat_vector 和 random_vectors_prj_dist_mat_vector 的皮尔逊相关系数
        corr, _ = pearsonr(random_vectors_dist_mat_vector, random_vectors_prj_dist_mat_vector)

        corrs.append(corr)

    plt.plot(corrs)
    plt.show()
    
